#!/Pub/Apps/Cellar/R/4.1.2/lib/R/bin/Rscript
suppressPackageStartupMessages(library(argparse))

parser <- ArgumentParser(
    prog = "inferCNV_py_plot_in_r.r",
    description = "对inferCNV_py结果简单可视化。",
    formatter_class = "argparse.RawTextHelpFormatter"
    # usage = "seurat2loom_exprMat.r -h; seurat2loom_exprMat.r seurat_obj od"
)

parser$add_argument("-od", "--output_dir", help = "结果输出路径。")
parser$add_argument("-k", "--key_column", help = "分群所在列。")
parser$add_argument("-r", "--reference", help = "对照组细胞类型。")

args <- parser$parse_args()

args.key_column <- args$key_column
args.reference <- args$reference
args.output_dir <- args$output_dir

suppressMessages(library(tidyverse))
tryCatch(
    {
        df <- arrow::read_feather(str_glue("{args.output_dir}/data_frame_meta.data.feather"))
    },
    error = function(e) {
        cli::cli_alert_danger("未能读取到data_frame_meta.data.feather，或者其他错误，vlnplot绘制失败。")
        q()
    }
)

n = length(unique(df[[args.key_column]]))
if(n > 10){
    cmap = colorRampPalette(ggsci::pal_jco()(10))(n)
}else {
   cmap = 'jco'
}

library(ggpubr)
plots <- ggpubr::ggviolin(df,
    x = args.key_column, y = "cnv_score",fill = args.key_column, title = args.key_column,
    palette = cmap, xlab = F, legend = "none"
) +
    theme(axis.text.x.bottom = element_text(angle = 90, vjust = .5, hjust = 1))

ggsave(str_glue("{args.output_dir}/infercnv_{args.key_column}_hclust.Vlnplot.pdf"), plots, width = 4.5, height = 4)

df2 <- data.frame(
    cell = df$cell,
    cnv_score = df$cnv_score,
    cls = df[[args.key_column]],
    type = ifelse(df[[args.key_column]] == args.reference, "reference", "non_ref"),
    x = 1
)

p <- tidyheatmaps::tidy_heatmap(df2, x, cell,color_legend_n = 11,
    values = cnv_score, annotation_col = c(cls, type), show_rownames = F,
    show_colnames = F, cluster_cols = T,
    clustering_distance_cols = "euclidean", silent = T, clustering_method = "ward.D2"
)

ggsave(str_glue("{args.output_dir}/infercnv_{args.key_column}_hclust.Heatplot.pdf"), p, width = 5, height = 3.75)
